from pathlib import Path

from ipyniivue import download_dataset

BASE_API_URL = "https://niivue.com/demos/images/"
DATA_FOLDER = Path("images")

# Download data for example
download_dataset(
    BASE_API_URL,
    DATA_FOLDER,
    files=[
        "CIT168.mz3",
        "BrainMesh_ICBM152.lh.mz3",
        "aal.mz3",
        "ColumnMajorOrder.gii",
        "mni_format.obj.gz",
        "simplify_brain.obj",
        "lh.pial",
        "sub-test02_left_hemisphere.srf.gz",
        "tract.SLF1_R.tck",
        "colby.trk",
        "colby.trx",
        "TR_S_R.tt.gz",
        "tract.FAT_R.vtk",
        "water-color.wrl",
        "MolView-sticks-color_38.x3d",
    ],
)

CIT168.mz3 already exists.
BrainMesh_ICBM152.lh.mz3 already exists.
aal.mz3 already exists.
ColumnMajorOrder.gii already exists.
mni_format.obj.gz already exists.
simplify_brain.obj already exists.
lh.pial already exists.
sub-test02_left_hemisphere.srf.gz already exists.
tract.SLF1_R.tck already exists.
colby.trk already exists.
colby.trx already exists.
TR_S_R.tt.gz already exists.
tract.FAT_R.vtk already exists.
water-color.wrl already exists.
MolView-sticks-color_38.x3d already exists.
Dataset downloaded successfully to images.

import ipywidgets as widgets
from IPython.display import display

import ipyniivue

nv = ipyniivue.NiiVue()
nv.opts.back_color = [0.9, 0.9, 1, 1]
nv.opts.show_3d_crosshair = True
nv.opts.slice_type = ipyniivue.SliceType.RENDER
nv.opts.show_legend = False

meshes_init = [
    {"path": DATA_FOLDER / "CIT168.mz3", "rgba255": [0, 0, 255, 255]},
    {"path": DATA_FOLDER / "BrainMesh_ICBM152.lh.mz3", "rgba255": [222, 164, 164, 255]},
]

nv.load_meshes(meshes_init)

nv.set_clip_plane(-0.1, 270, 0)


@nv.on_mesh_loaded
def on_init_load(mesh):
    """Set mesh shader initial value on load."""
    if len(nv.meshes) > 1:
        nv.set_mesh_shader(nv.meshes[1].id, "Outline")
        nv.on_mesh_loaded(on_init_load, remove=True)


# Cell 2: UI Controls and Logic

# Color Slider
color_slider = widgets.IntSlider(value=128, min=0, max=255, description="Color")


def on_color_change(change):
    """Handle color change."""
    if not nv.meshes:
        return
    val = change["new"]
    target_mesh = nv.meshes[-1]
    nv.set_mesh_property(target_mesh.id, "rgba255", [val, 164, 164, 255])


color_slider.observe(on_color_change, names="value")

# Reverse Faces Button
reverse_btn = widgets.Button(description="Reverse Faces")


def on_reverse_click(b):
    """Reverse mesh faces."""
    if not nv.meshes:
        return
    if len(nv.meshes) > 0:
        nv.reverse_faces(nv.meshes[0].id)
    if len(nv.meshes) > 1:
        nv.reverse_faces(nv.meshes[1].id)


reverse_btn.on_click(on_reverse_click)

# Alpha Select
alpha_dropdown = widgets.Dropdown(
    options=[("Transparent", 0), ("Translucent", 1), ("Opaque", 2)],
    value=2,
    description="Alpha",
)


def on_alpha_change(change):
    """Handle alpha change."""
    if not nv.meshes:
        return
    idx = change["new"]
    target_mesh = nv.meshes[-1]

    nv.set_mesh_property(target_mesh.id, "visible", idx > 0)

    if idx == 1:
        nv.set_mesh_property(target_mesh.id, "opacity", 0.2)
    elif idx == 2:
        nv.set_mesh_property(target_mesh.id, "opacity", 1.0)
    else:
        nv.set_mesh_property(target_mesh.id, "opacity", 0.0)

    nv.draw_scene()


alpha_dropdown.observe(on_alpha_change, names="value")

# Crosshair Checkbox
cross_check = widgets.Checkbox(value=True, description="Crosshairs")


def on_cross_change(change):
    """Handle crosshair checkbox."""
    nv.opts.show_3d_crosshair = change["new"]
    nv.draw_scene()


cross_check.observe(on_cross_change, names="value")

# Shader Controls
shader_names = nv.mesh_shader_names()

shader_last = widgets.Dropdown(
    value="Outline", options=shader_names, description="Shader"
)


def change_shader(change, mesh_index):
    """Handle shader change."""
    if not nv.meshes or len(nv.meshes) <= mesh_index:
        if mesh_index == -1 and len(nv.meshes) > 0:
            target_mesh = nv.meshes[-1]
        else:
            return
    else:
        target_mesh = nv.meshes[mesh_index]

    nv.set_mesh_shader(target_mesh.id, change["new"])


shader_last.observe(lambda c: change_shader(c, -1), names="value")

# Mesh Loading Buttons
mesh_files = [
    "aal.mz3",
    "ColumnMajorOrder.gii",
    "mni_format.obj.gz",
    "simplify_brain.obj",
    "lh.pial",
    "sub-test02_left_hemisphere.srf.gz",
    "tract.SLF1_R.tck",
    "colby.trk",
    "colby.trx",
    "TR_S_R.tt.gz",
    "tract.FAT_R.vtk",
    "water-color.wrl",
    "MolView-sticks-color_38.x3d",
]

mesh_dropdown = widgets.Dropdown(options=mesh_files, description="Load Mesh")
load_btn = widgets.Button(description="Load")


def on_load_click(b):
    """Handle load click."""
    filename = mesh_dropdown.value
    nv.load_meshes([{"path": DATA_FOLDER / filename}])


load_btn.on_click(on_load_click)

header = widgets.HBox([color_slider, reverse_btn, alpha_dropdown, cross_check])
footer = widgets.VBox(
    [
        shader_last,
        widgets.HBox([widgets.Label("Additional Meshes:"), mesh_dropdown, load_btn]),
    ]
)

display(header, nv, footer)